#!/usr/bin/python
import sys 
import subprocess 
import shlex 
import re 
import os

defaults = {"e_threshold":0.005,"cutoff":0.7,"min_tms":2,"max_gi":200,"max_gsat":10}

handles = []
status = []
valid_acc = []

def protocol1(acc):
	globalstatus = get_global_status(acc)
	if globalstatus is "DONE" or globalstatus is "TMS" or globalstatus is "BROKEN" :
		return
	cmd = "php /usr/bin/protocol1 -q 1 -a %s -e %f -o %s -c %f -x %i > /dev/null" %(acc,defaults["e_threshold"],"P1_"+acc,defaults["cutoff"],defaults['max_gi'])
	args = shlex.split(cmd)
	return subprocess.Popen(args)
	
def protocol2():
	cmd = 'php /usr/bin/protocol2 -q 1 -s %s -t %s -o %s -m %i > /dev/null' %(sys.argv[1]+'.faa',sys.argv[2]+'.faa','./',defaults['max_gsat'])
	args = shlex.split(cmd)
	return subprocess.Popen(args)
	
def getstatus(acc):
	if os.path.isfile("./P1_"+acc+"/status") is not True:
		return 0
	status = open("./P1_"+acc+"/status","r");
	status = status.read()
	if len(status) > 0 :
		return int(status == "DONE")
	else:
		return 0
		
def run_hmmtop(acc):
	for i in range(len(acc)):
		if os.path.isfile(acc[i]+".faa") is True:
			continue
		file = open("./P1_"+acc[i]+"/sequences_fasta.xml.faa")
		fasta = file.read()
		fasta = filter(lambda x:len(x)>0,fasta.split(">"))
		fasta = dict(zip(range(len(fasta)),fasta))
		cmd = "hmmtop -if=%s sf=FAS -is=pseudo -pi=spred -of=/tmp/ramdisk/%s.hmmtop > /dev/null" %("./P1_"+acc[i]+"/sequences_fasta.xml.faa",acc[i])
		args = shlex.split(cmd)
		handles.append(subprocess.Popen(args))
	for h in handles:
		h.wait()
	return
		
def tms_filter(acc):
	tmss=[]
	file = open("./P1_"+acc+"/sequences_fasta.xml.faa")
	fasta = file.read()
	fasta = filter(lambda x:len(x)>0,fasta.split(">"))
	fasta = dict(zip(range(len(fasta)),fasta))
	results = open("/tmp/ramdisk/"+acc+".hmmtop")
	results = filter(lambda x:len(x)>0,results.read().split(">"))
	for x in results:
		tms = re.search('(IN|OUT)\s+(\d{1,3})',x)
		tmss.append(tms.groups()[1])
	for i in range(len(tmss)):
		if int(tmss[i]) < defaults['min_tms']:
			fasta.pop(i)
	return fasta.values()

def cleanup():
	for acc in sys.argv[1:]:
		try:
			os.remove("/tmp/ramdisk/"+acc+".hmmtop")
		except:
			pass
	return

def send_acc(acc):
	statuspath = "./P1_"+acc+"/status"
	filepath = "./%s.faa" %(acc)
	if os.path.isfile(statuspath) is False:
		return # No status file to send, nothing to do
	status = open(filepath,"r")
	localstatus = status.read().strip()
	globalstatus = get_global_status(acc)
	if globalstatus is "DONE" or globalstatus is "BROKEN" or globalstatus is "TMS":
		return
	else:
		mysql("insert into acc ( datetime, acc) values ( NOW(), '"+acc+"')")
		if localstatus is "DONE":
			if os.path.isfile(filepath) is True:
				# + Update DONE & Upload
				send_file("./"+acc+".faa","/Library/webserver/acc/"+acc+".faa")
				mysql("update acc set status='DONE' where acc='"+acc+"'")
			else:
				# + Update TMS
				mysql("update acc set status='TMS' where acc='"+acc+"'")
		else:
			# + Update status'
			mysql("update acc set status='"+localstatus+"' where acc='"+acc+"'")
	return
	
def get_acc(acc):
	globalstatus = get_global_status(acc)
	if globalstatus is "DONE":
		get_file("/Library/webserver/acc/"+acc+".faa","./"+acc+".faa")
		return True
	return False
	
def mysql(query): # LOL
	q="php ./db.php \"%s\" 0"%(query)
	db = subprocess.Popen(q, shell="True", stdout=subprocess.PIPE)
	result=db.communicate()[0]
	try:
		return int(result)
	except:
		if len(result) > 0:
			return result
		return False
		
def get_global_status(acc):
	return status = mysql("Select status from acc where acc='"+acc+"'")

def main():
	for acc in sys.argv[1:]:
		get_acc(acc)
		if os.path.isfile(acc+".faa") is True:
			continue
		if getstatus(acc) is not 1:
			handles.append(protocol1(acc))
	for handle in handles:
		try:
			handle.wait()
		except:
			pass
	if sum([getstatus(x) for x in sys.argv[1:]]) is not 2 and os.path.isfile(sys.argv[1]+".faa") is False and os.path.isfile(sys.argv[2]+".faa") is False:
		# Protocol 1 failed
		sys.exit()
	# We meow have two working FASTA Files. Now we check for good TMS scores.
	run_hmmtop(sys.argv[1:])
	for acc in sys.argv[1:]:
		if os.path.isfile(acc+".faa") is True:
			valid_acc.append(80) # Error code, file already exists
			continue
		fasta=tms_filter(acc)
		if len(fasta)>0 :
			valid_acc.append(acc)
			handle = open(acc+".faa","w")
			handle.write(">"+">".join(fasta))
	if len(valid_acc) is not 2 :
		sys.exit() 
		# One or more ACCs failed the TMS test.
	cleanup()
	# Protocol1 success, now for P2
	P2 = protocol2()
	P2.wait()
	
	
if __name__ == "__main__":
	main()